#include <cmath>
#include <iostream>
#include <random>
#include <vector>

#define _USE_MATH_DEFINES
#include <math.h>

#include <imgui.h>

#include <backends/imgui_impl_sdl3.h>
#include <backends/imgui_impl_sdlrenderer3.h>

#include <SDL3/SDL.h>

// --- Configuration Constants ---
const int SCREEN_WIDTH = 1000;
const int SCREEN_HEIGHT = 800;
const int NUM_PARTICLES = 5000;
const char *WINDOW_TITLE = "SDL3 2D Fluid Simulation with ImGui";

// --- Particle Structure ---
struct Particle {
  float x = 0.0f, y = 0.0f;   // Position
  float vx = 0.0f, vy = 0.0f; // Velocity
  float fx = 0.0f, fy = 0.0f; // Force
  float rho = 0.0f, p = 0.0f; // Density and Pressure
};

// --- SPH Simulation Constants (Now modifiable via ImGui) ---
struct SimConstants {
  float G = 2500.0f;         // Gravity
  float REST_DENS = 1000.0f; // Rest density
  float GAS_CONST = 2000.0f; // Gas constant
  float H = 16.0f;           // Smoothing kernel radius
  float MASS = 65.0f;        // Particle mass
  float VISC = 250.0f;       // Viscosity constant
  float DT = 0.0008f;        // Integration timestep
};

// --- Simulation Class ---
class FluidSimulator {
public:
  // Pass constants by pointer to allow modification from outside
  FluidSimulator(SimConstants *constants) : m_constants(constants) { reset(); }

  void reset() {
    m_particles.clear();
    m_particles.resize(NUM_PARTICLES);
    // Place particles in a grid
    int particles_per_row = static_cast<int>(sqrt(NUM_PARTICLES));
    for (int i = 0; i < NUM_PARTICLES; ++i) {
      m_particles[i].x = 100 + (i % particles_per_row) * m_constants->H * 0.5f;
      m_particles[i].y = 100 + (i / particles_per_row) * m_constants->H * 0.5f;
      m_particles[i].vx = m_particles[i].vy = 0.0f;
    }
  }

  // Main update function for the physics simulation
  void update(float mouseX, float mouseY, bool mouseDown) {
    calculateDensityAndPressure();
    calculateForces();
    integrate(mouseX, mouseY, mouseDown);
  }

  const std::vector<Particle> &getParticles() const { return m_particles; }

private:
  SimConstants *m_constants;
  std::vector<Particle> m_particles;

  void calculateDensityAndPressure() {
    const float hsq = m_constants->H * m_constants->H;
    const float poly6 = 315.0f / (64.0f * M_PI * powf(m_constants->H, 9));

#pragma omp parallel for schedule(static)
    for (int i = 0; i < m_particles.size(); ++i) {
      Particle &pi = m_particles[i];
      pi.rho = 0.0f;
      for (Particle &pj : m_particles) {
        float dx = pj.x - pi.x;
        float dy = pj.y - pi.y;
        float r2 = dx * dx + dy * dy;

        if (r2 < hsq) {
          pi.rho += m_constants->MASS * poly6 * powf(hsq - r2, 3);
        }
      }
      pi.p = m_constants->GAS_CONST * (pi.rho - m_constants->REST_DENS);
    }
  }

  void calculateForces() {
    const float spiky_grad = -45.0f / (M_PI * powf(m_constants->H, 6));
    const float visc_lap = 45.0f / (M_PI * powf(m_constants->H, 6));
    const float h = m_constants->H;
    const float hsq = h * h;

#pragma omp parallel for schedule(static)
    for (int i = 0; i < m_particles.size(); ++i) {
      Particle &pi = m_particles[i];
      float f_press_x = 0.0f, f_press_y = 0.0f;
      float f_visc_x = 0.0f, f_visc_y = 0.0f;

      for (Particle &pj : m_particles) {
        if (&pi == &pj)
          continue;

        float dx = pj.x - pi.x;
        float dy = pj.y - pi.y;
        float r2 = dx * dx + dy * dy;

        if (r2 < hsq) {
          float r = sqrtf(r2);
          if (r < 1e-6f)
            r = 1e-6f; // Avoid division by zero

          if (pj.rho > 1e-6f) {
            // Pressure force
            float C = (pi.p + pj.p) / (2.0f * pj.rho) * spiky_grad * powf(h - r, 2);
            f_press_x += C * dx / r;
            f_press_y += C * dy / r;

            // Viscosity force
            f_visc_x += m_constants->VISC * (pj.vx - pi.vx) / pj.rho * visc_lap * (h - r);
            f_visc_y += m_constants->VISC * (pj.vy - pi.vy) / pj.rho * visc_lap * (h - r);
          }
        }
      }

      float f_grav_y = m_constants->G;
      pi.fx = f_press_x + f_visc_x;
      pi.fy = f_press_y + f_visc_y + f_grav_y;
    }
  }

  void integrate(float mouseX, float mouseY, bool mouseDown) {
    const float bound_damping = -0.5f;
    const float view_width = static_cast<float>(SCREEN_WIDTH);
    const float view_height = static_cast<float>(SCREEN_HEIGHT);
    const float h = m_constants->H;

#pragma omp parallel for schedule(static)
    for (int i = 0; i < m_particles.size(); ++i) {
      Particle &p = m_particles[i];
      // Add mouse interaction force if not interacting with UI
      if (mouseDown && !ImGui::GetIO().WantCaptureMouse) {
        float dx = p.x - mouseX;
        float dy = p.y - mouseY;
        float r2 = dx * dx + dy * dy;
        if (r2 < 10000) { // Mouse radius of 100px
          p.fx -= dx * 100.0f;
          p.fy -= dy * 100.0f;
        }
      }

      // Forward Euler integration
      if (p.rho > 1e-6f) {
        p.vx += p.fx / p.rho * m_constants->DT;
        p.vy += p.fy / p.rho * m_constants->DT;
      }
      p.x += p.vx * m_constants->DT;
      p.y += p.vy * m_constants->DT;

      // Enforce boundary conditions
      if (p.x - h < 0.0f) {
        p.vx *= bound_damping;
        p.x = h;
      }
      if (p.x + h > view_width) {
        p.vx *= bound_damping;
        p.x = view_width - h;
      }
      if (p.y - h < 0.0f) {
        p.vy *= bound_damping;
        p.y = h;
      }
      if (p.y + h > view_height) {
        p.vy *= bound_damping;
        p.y = view_height - h;
      }
    }
  }
};

// --- Main Application ---
int main(int argc, char *argv[]) {
  // Initialize SDL
  if (SDL_Init(SDL_INIT_VIDEO) < 0) {
    std::cerr << "Could not initialize SDL: " << SDL_GetError() << std::endl;
    return 1;
  }

  // Create Window and Renderer
  SDL_Window *window = nullptr;
  SDL_Renderer *renderer = nullptr;
  if (SDL_CreateWindowAndRenderer(WINDOW_TITLE, SCREEN_WIDTH, SCREEN_HEIGHT, SDL_WINDOW_RESIZABLE, &window, &renderer) < 0) {
    std::cerr << "Could not create window and renderer: " << SDL_GetError() << std::endl;
    SDL_Quit();
    return 1;
  }

  // --- ImGui Setup ---
  IMGUI_CHECKVERSION();
  ImGui::CreateContext();
  ImGuiIO &io = ImGui::GetIO();
  (void)io;
  ImGui::StyleColorsDark();
  ImGui_ImplSDL3_InitForSDLRenderer(window, renderer);
  ImGui_ImplSDLRenderer3_Init(renderer);

  // Create simulation constants and simulator instance
  SimConstants constants;
  FluidSimulator sim(&constants);

  bool running = true;
  SDL_Event event;

  float mouseX = 0.0f, mouseY = 0.0f;
  bool mouseDown = false;

  // Main loop
  while (running) {
    // Event handling
    while (SDL_PollEvent(&event)) {
      ImGui_ImplSDL3_ProcessEvent(&event);
      if (event.type == SDL_EVENT_QUIT) {
        running = false;
      }
      if (event.type == SDL_EVENT_MOUSE_BUTTON_DOWN && event.button.button == SDL_BUTTON_LEFT) {
        mouseDown = true;
      }
      if (event.type == SDL_EVENT_MOUSE_BUTTON_UP && event.button.button == SDL_BUTTON_LEFT) {
        mouseDown = false;
      }
      if (event.type == SDL_EVENT_MOUSE_MOTION) {
        SDL_GetMouseState(&mouseX, &mouseY);
      }
    }

    // --- ImGui Frame Start ---
    ImGui_ImplSDLRenderer3_NewFrame();
    ImGui_ImplSDL3_NewFrame();
    ImGui::NewFrame();

    // --- ImGui UI Definition ---
    {
      ImGui::Begin("SPH Settings");

      ImGui::Text("Application average %.3f ms/frame (%.1f FPS)", 1000.0f / ImGui::GetIO().Framerate,
                  ImGui::GetIO().Framerate);
      ImGui::Separator();

      ImGui::SliderFloat("Gravity", &constants.G, 0.0f, 10000.0f);
      ImGui::SliderFloat("Viscosity", &constants.VISC, 0.0f, 1000.0f);
      ImGui::SliderFloat("Gas Constant", &constants.GAS_CONST, 500.0f, 5000.0f);
      ImGui::Separator();
      ImGui::SliderFloat("Smoothing Radius (H)", &constants.H, 5.0f, 50.0f);
      ImGui::SliderFloat("Particle Mass", &constants.MASS, 10.0f, 200.0f);

      if (ImGui::Button("Reset Simulation")) {
        sim.reset();
      }

      ImGui::End();
    }

    // Update simulation
    sim.update(mouseX, mouseY, mouseDown);

    // --- Rendering ---
    ImGui::Render();
    SDL_SetRenderDrawColor(renderer, 20, 25, 30, 255);
    SDL_RenderClear(renderer);

    // Draw particles
    const auto &particles = sim.getParticles();
    SDL_SetRenderDrawColor(renderer, 70, 150, 255, 255);
    for (const auto &p : particles) {
      SDL_FRect rect = {p.x - 2.0f, p.y - 2.0f, 4.0f, 4.0f};
      SDL_RenderFillRect(renderer, &rect);
    }

    // Render ImGui draw data
    ImGui_ImplSDLRenderer3_RenderDrawData(ImGui::GetDrawData(), renderer);

    // Present the frame
    SDL_RenderPresent(renderer);
  }

  // --- Cleanup ---
  ImGui_ImplSDLRenderer3_Shutdown();
  ImGui_ImplSDL3_Shutdown();
  ImGui::DestroyContext();

  SDL_DestroyRenderer(renderer);
  SDL_DestroyWindow(window);
  SDL_Quit();

  return 0;
}
